# PYTHON_ARGCOMPLETE_OK
# 
# MIT License
# 
# Copyright (c) 2024 Jianbo ZHU
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import os
import logging
import argparse

import argcomplete

from .environ import *
from . import core


def supercell(args: argparse.Namespace):
    if nref := args.cubicize:
        posdict = core.read_poscar('POSCAR', skip_coordinates=True)
        dimmat = core.cubicize_dimmat(
            lattice=posdict["lattice"]*posdict["factor"],
            natom=sum(posdict["numbers"]),
            nref=nref,
            nmin=args.nmin,
            nmax=args.nmax,
        )
        core.write_dimmat(dimmat, args.abort_on_exist)
    elif args.dim:
        core.write_dimmat(args.dim, args.abort_on_exist)

    dimstr = core.read_dimmat()
    core.cmd(f"{args.exec_phonopy} -d --dim='{dimstr}'")

def _options_supercell(parser_obj):
    parser_obj.add_argument("--dim", metavar="DIMMAT", help=f"Supercell dimensions (default: read from {PHONOKIT_DIMMAT})")
    parser_obj.add_argument("--cubicize", metavar='NREF', type=int, help="Build a near cubic supercell based on a reference atom count")
    parser_obj.add_argument("--nmin", type=int, help="Limit the minimum number of atoms allowed in the supercell")
    parser_obj.add_argument("--nmax", type=int, help="Limit the maximum number of atoms allowed in the supercell")
    return supercell

def prepare(args: argparse.Namespace, scaling=None):
    if (scaling is None) and (args.scaling is not None):
        for s in args.scaling.split():
            prepare(args, scaling=float(s))
        return

    if (args.ext is None) and (scaling is not None):
        fext = f"{round(scaling*1E4):05d}"
    else:
        fext = args.ext

    fdir = core.cleandir(args.caldir, args.abort_on_exist, ext=fext)
    vaspkit_kp = f"{args.exec_vaspkit} -task 102 -kpr {args.kpr} -kps {args.kps.upper()[0]}"

    if any([args.dfpt, args.born, args.scf]):
        if args.dfpt:
            sposcar = core.checkdir(args.sposcar or "SPOSCAR")
        else:
            sposcar = core.checkdir(args.sposcar or "POSCAR")
        with core.workdir(fdir) as orig:
            core.cmd(f"cp {orig}/{sposcar} POSCAR")
            if args.kpoints and args.potcar:
                logging.debug("Skipping invoking vaspkit (INCAR, KPOINTS, POTCAR are provided)")
            else:
                core.cmd(vaspkit_kp)

            if (incar := args.incar):
                logging.info(f"Using custom INCAR file:  {incar}")
                core.cmd(f"cp {orig}/{incar} INCAR")
            elif args.dfpt:
                logging.info("Generating the default INCAR file for DFPT calculations...")
                with open("INCAR", "w") as f:
                    f.write(PHONOKIT_INCAR_DFPT)
            elif args.born:
                logging.info("Generating the default INCAR file for BORN charge calculations...")
                with open("INCAR", "w") as f:
                    f.write(PHONOKIT_INCAR_BORN)
            elif args.scf:
                logging.info("Generating the default INCAR file for SCF calculations...")
                with open("INCAR", "w") as f:
                    f.write(PHONOKIT_INCAR_SCF.format(args.ncore))
            else:
                raise RuntimeError("Unexpected error: Failed to determine the calculation type.")

            if (kpoints := args.kpoints):
                logging.info(f"Using custom KPOINTS file:  {kpoints}")
                core.cmd(f"cp {orig}/{kpoints} KPOINTS")
            if (potcar := args.potcar):
                logging.info(f"Using custom POTCAR file:  {potcar}")
                core.cmd(f"cp {orig}/{potcar} POTCAR")

            if scaling:
                logging.info(f"Overriding the scaling factor of the POSCAR file:  {scaling}")
                core.cmd(f"sed -i '2c\   {scaling:.6f}' POSCAR")
        return

    # finite-displacement workflow
    sposcar = core.checkdir(args.sposcar or "SPOSCAR")
    core.cmd(f"mkdir -p {fdir}/inputf")
    with core.workdir(f"{fdir}/inputf") as orig:
        core.cmd(f"cp {orig}/{sposcar} POSCAR")
        if args.kpoints and args.potcar:
            logging.debug("Skipping invoking vaspkit (INCAR, KPOINTS, POTCAR are provided)")
        else:
            core.cmd(vaspkit_kp)

        if (incar := args.incar):
            logging.info(f"Using custom INCAR file:  {incar}")
            core.cmd(f"cp {orig}/{incar} INCAR")
        else:
            logging.info(f"Generating the default INCAR file with NCORE = {args.ncore}...")
            with open("INCAR", "w") as f:
                f.write(PHONOKIT_INCAR_SCF.format(args.ncore))

        if (kpoints := args.kpoints):
            logging.info(f"Using custom KPOINTS file:  {kpoints}")
            core.cmd(f"cp {orig}/{kpoints} KPOINTS")
        if (potcar := args.potcar):
            logging.info(f"Using custom POTCAR file:  {potcar}")
            core.cmd(f"cp {orig}/{potcar} POTCAR")

    for uid in range(1, 1000):
        if not os.path.exists(fpos := f"POSCAR-{uid:03g}"):
            break
        fuid = f"FD-{uid:03g}"
        core.cmd(f"cp -r {fdir}/inputf {fdir}/{fuid} && cp {fpos} {fdir}/{fuid}/POSCAR")
    else:
        raise RuntimeError("Unexpected error: Number of displacement files exceeds 999.")

    if scaling:
        logging.info(f"Overriding the scaling factor of all POSCAR files:  {scaling}")
        core.cmd(f"sed -i '2c\   {scaling:.6f}' {fdir}/FD-*/POSCAR")

def _options_prepare(parser_obj):
    parser_obj.add_argument("--dfpt", action="store_true", help="Use DFPT instead of finite-displacements (FD) workflow.")
    parser_obj.add_argument("--born", action="store_true", help="Prepare VASP input files for BORN charge calculations. <experimental>")
    parser_obj.add_argument("--scf", action="store_true", help="Prepare VASP input files for self-consistent-field calculations. <experimental>")
    parser_obj.add_argument("--caldir", default=PHONOKIT_CALDIR, help=f"Output directory (default: {PHONOKIT_CALDIR})")
    parser_obj.add_argument("-e", "--ext", help="Assign the extension (default: SCALING_FACTOR)")
    parser_obj.add_argument("--sposcar", help="Specify the input POSCAR file (default: POSCAR or SPOSCAR)")
    parser_obj.add_argument("--scaling", metavar='FACTOR_LIST', help="Override the scaling factor of the SPOSCAR file (default: UNSET)")
    parser_obj.add_argument("--kps", default="Gamma", help="K-mesh scheme: Gamma[default] or Monkhorst-Pack")
    parser_obj.add_argument("--kpr", type=float, default=0.03, help="K-spacing value in unit of 2*pi/Angstrom (default: 0.03)")
    parser_obj.add_argument("--kpoints", help="Assign a custom KPOINTS file to override the default (default: UNSET)")
    parser_obj.add_argument("--potcar", help="Assign a custom POTCAR file to override the default (default: UNSET)")
    parser_obj.add_argument("--incar", help="Assign a custom INCAR file to override the default (default: UNSET)")
    parser_obj.add_argument("--ncore", type=int, default=PHONOKIT_VASP_NCORE, help=f"Parameter NCORE in INCAR (default: {PHONOKIT_VASP_NCORE})")
    return prepare

def collect(args: argparse.Namespace):
    cdir = core.checkdir(args.caldir, ext=args.ext)
    core.cmd(f"cp POSCAR {PHONOKIT_DISPYAML} {cdir}")
    if args.dfpt:
        force_type = "FORCE_CONSTANTS"
        with core.workdir(cdir):
            core.checkdir("vasprun.xml", abort_msg="No VASP output vasprun.xml found.")
            core.cmd(f"{args.exec_phonopy} --fc vasprun.xml")
    else:
        force_type = "FORCE_SETS"
        with core.workdir(cdir):
            if not os.path.exists("FD-001"):
                err_msg = "No finite-displacement calculation results found."
                if os.path.exists("vasprun.xml"):
                    err_msg += " You might try using --dfpt."
                raise RuntimeError(err_msg)
            core.cmd(f"{args.exec_phonopy} -f FD-*/vasprun.xml")
    logging.info(f"Collect IFC results ( {force_type} ) scuccessfully.")

    if (borndir := args.with_born):
        bdir = core.checkdir(borndir)
        with core.workdir(bdir):
            logging.info(f"Collecting BORN charge results from {bdir}...")
            core.cmd(f"phonopy-vasp-born | tee BORN")
        bornfile = f"{bdir}/BORN"
    else:
        bornfile = ""

    fdir = core.cleandir(args.forcedir, args.abort_on_exist, ext=args.ext)
    core.cmd(f"cp POSCAR {PHONOKIT_DISPYAML} {cdir}/{force_type} {bornfile} {fdir}")
    with core.workdir(fdir):
        dimstr = core.read_dimmat()
        core.write_dimmat(dimstr)
        logging.debug(f"Write to supercell dimensions to {PHONOKIT_DIMMAT}")


def _options_collect(parser_obj):
    parser_obj.add_argument("--dfpt", action="store_true", help="Use DFPT instead of finite-displacements (FD) workflow.")
    parser_obj.add_argument("--caldir", default=PHONOKIT_CALDIR, help=f"Calculation directory (default: {PHONOKIT_CALDIR})")
    parser_obj.add_argument("--forcedir", default=PHONOKIT_FORCE, help=f"Output directory (default: {PHONOKIT_FORCE})")
    parser_obj.add_argument("-e", "--ext", help="Assign the extension (default: UNSET)")
    parser_obj.add_argument("--with-born", metavar="BORN_DIR", help="Collect BORN charge results (default: UNSET)")
    return collect

def band(args: argparse.Namespace):
    if band_yaml := args.yaml:
        _band_yaml(band_yaml, args)
        return
    fdir = core.checkdir(args.forcedir, ext=args.ext)
    odir = core.cleandir(args.outdir, args.abort_on_exist, copyfrom=fdir, ext=args.ext)
    with core.workdir(odir):
        if args.band is None:
            band, band_labels = core.get_kpath("POSCAR", args.exec_vaspkit)
        else:
            band = args.band
            band_labels = args.band_labels
        cmdstr = f"{args.exec_phonopy}"
        cmdstr += f" --dim='{core.read_dimmat()}'"
        cmdstr += f" --mesh='{args.mesh}'"
        cmdstr += " -p -s" if args.plot else ""
        cmdstr += " --nac" if args.nac else ""
        cmdstr += " --readfc" if core.read_forcetype() == "FORCE_CONSTANTS" else ""
        cmdstr += f" --band='{band}'"
        cmdstr += f" --band-labels='{band_labels}'"
        cmdstr += f" --band-points={args.band_points}" if args.band_points is not None else ""
        cmdstr += " --band-connection" if args.band_connection else ""
        cmdstr += f" --pdos='{args.pdos}'" if args.pdos is not None else ""
        cmdstr += f" --sigma={args.sigma}" if args.sigma is not None else ""
        core.cmd(cmdstr)

        if os.path.exists("projected_dos.dat"):
            logging.debug("Rename projected_dos.dat to projected_dos.txt...")
            core.cmd("mv projected_dos.dat projected_dos.txt")
        if os.path.exists("total_dos.dat"):
            logging.debug("Rename total_dos.dat to total_dos.txt...")
            core.cmd("mv total_dos.dat total_dos.txt")
        _band_yaml("band.yaml", args)

def _band_yaml(bandfile, args: argparse.Namespace):
    logging.debug("Write band structure to file:  band.txt...")
    core.parse_band(bandfile, "band.txt")
    # core.cmd('echo -n "#  kpath  branch-1 branch-2 ..." > band.txt')
    # core.cmd(r"""awk '/distance/{printf "\n%16.10f", $2} /frequency/{printf "%16.10f", $2} END{print "" }' band.yaml >> band.txt""")
    if args.skip_klabels:
        logging.debug("Skip parseing k-klables.")
    else:
        logging.debug("Write k-path labels to file:  band_labels.txt")
        core.parse_klabeles("band.yaml", "band_labels.txt", args.legacy)

def _options_band(parser_obj):
    parser_obj.add_argument("-y", "--yaml", metavar="BAND.YAML", help="Start from a yaml file rather than force constants.")
    parser_obj.add_argument("--plot", action="store_true", help="Plot the phonon band structureand save")
    parser_obj.add_argument("--nac", action="store_true", help="Including non-analytical term correction (NAC).")
    parser_obj.add_argument("--forcedir", default=PHONOKIT_FORCE, help=f"Directory of IFC results (default: {PHONOKIT_FORCE})")
    parser_obj.add_argument("--outdir", default="band", help="Output directory (default: band)")
    parser_obj.add_argument("-e", "--ext", help="Assign the extension (default: UNSET)")
    parser_obj.add_argument("--mesh", metavar="'N_a N_b N_c'", default="32 32 32", help="Sampling mesh (default: '32 32 32')")
    parser_obj.add_argument("--band", metavar="'K_a1 K_b1 K_c1 K_a2 K_b2 K_c2 ...'", help="Sampling band paths (default: from VASPKIT)")
    parser_obj.add_argument("--band-labels", metavar="'label1 label2 ...'", help="Labels for band paths (default: from VASPKIT)")
    parser_obj.add_argument("--band-points", type=int, help="Number of sampling points for path ends (default: 51)")
    parser_obj.add_argument("--band-connection", action="store_true", help="Try to obtain better band connections")
    parser_obj.add_argument("--pdos", help="Calculate projected DOS instaed of the total (default: UNSET)")
    parser_obj.add_argument("--sigma", type=float, help="A smearing method is used instead of a linear tetrahedron method. (default: UNSET)")
    parser_obj.add_argument("--skip-klabels", action="store_true", help="Skip parseing k-klables.")
    parser_obj.add_argument("--legacy", action="store_true", help="Use legacy version.")
    return band

def tprop(args: argparse.Namespace):
    fdir = core.checkdir(args.forcedir, ext=args.ext)
    odir = core.cleandir(args.outdir, args.abort_on_exist, copyfrom=fdir, ext=args.ext)
    with core.workdir(odir):
        cmdstr = f"{args.exec_phonopy} -t"
        cmdstr += f" --dim='{core.read_dimmat()}'"
        cmdstr += f" --mesh='{args.mesh}'"
        cmdstr += " -p -s" if args.plot else ""
        cmdstr += " --nac" if args.nac else ""
        cmdstr += " --readfc" if core.read_forcetype() == "FORCE_CONSTANTS" else ""
        cmdstr += f" --tmin={args.tmin}" if args.tmin is not None else ""
        cmdstr += f" --tmax={args.tmax}" if args.tmax is not None else ""
        cmdstr += f" --tstep={args.tstep}" if args.tstep is not None else ""
        cmdstr += f" --cutoff-freq={args.cutoff_freq}" if args.cutoff_freq is not None else ""
        cmdstr += f" | tee {PHONOKIT_TMPFILE}"
        core.cmd(cmdstr)
        core.cmd(f"sed -n '/^#/,/^$/p' {PHONOKIT_TMPFILE} > tprop.txt && rm {PHONOKIT_TMPFILE}")
        logging.info("Write thermal properties to file:  tprop.txt")

def _options_tprop(parser_obj):
    parser_obj.add_argument("--plot", action="store_true", help="Plot the thermal properties and save")
    parser_obj.add_argument("--nac", action="store_true", help="Including non-analytical term correction (NAC).")
    parser_obj.add_argument("--forcedir", default=PHONOKIT_FORCE, help=f"Directory of IFC results (default: {PHONOKIT_FORCE})")
    parser_obj.add_argument("--outdir", default="tprop", help="Output directory (default: tprop)")
    parser_obj.add_argument("-e", "--ext", help="Assign the extension (default: UNSET)")
    parser_obj.add_argument("--mesh", default="32 32 32", help="Sampling mesh (default: 32 32 32)")
    parser_obj.add_argument("--tmin", type=float, help="Minimum calculated temperature (default: 0)")
    parser_obj.add_argument("--tmax", type=float, help="Maximum calculated temperature (default: 1000)")
    parser_obj.add_argument("--tstep", type=float, help="Calculated temperature step (default: 10)")
    parser_obj.add_argument("--cutoff-freq", type=float, help="Cutoff frequency (default: 0)")
    return tprop

def group_velocity(args: argparse.Namespace):
    if mesh_yaml := args.yaml:
        logging.debug(f"Write group velocities to file:  vg.txt")
        core.parse_vg(meshfile=mesh_yaml, output="vg.txt", with_q=args.q_position)
        return

    fdir = core.checkdir(args.forcedir, ext=args.ext)
    odir = core.cleandir(args.outdir, args.abort_on_exist, copyfrom=fdir, ext=args.ext)
    with core.workdir(odir):
        cmdstr = f"{args.exec_phonopy} --gv"
        cmdstr += f" --dim='{core.read_dimmat()}'"
        cmdstr += f" --mesh='{args.mesh}'"
        cmdstr += " --gc" if args.gc else ""
        cmdstr += " --nac" if args.nac else ""
        cmdstr += " --readfc" if core.read_forcetype() == "FORCE_CONSTANTS" else ""
        core.cmd(cmdstr)
        core.parse_vg(meshfile="mesh.yaml", output="vg.txt", with_q=args.q_position)
        logging.debug("Write group velocities to file:  vg.txt")

def _options_vg(parser_obj):
    parser_obj.add_argument("-y", "--yaml", metavar="MESH.YAML", help="Start from a yaml file rather than force constants.")
    parser_obj.add_argument("--nac", action="store_true", help="Including non-analytical term correction (NAC).")
    parser_obj.add_argument("--forcedir", default=PHONOKIT_FORCE, help=f"Directory of IFC results (default: {PHONOKIT_FORCE})")
    parser_obj.add_argument("--outdir", default="vg", help="Output directory (default: vg)")
    parser_obj.add_argument("-e", "--ext", help="Assign the extension (default: UNSET)")
    parser_obj.add_argument("--mesh", default="32 32 32", help="Sampling mesh (default: 32 32 32)")
    parser_obj.add_argument("--gc", "--gamma-center", action="store_true", help="Set mesh as Gamma center.")
    parser_obj.add_argument("--q-position", action="store_true", help="Write q-position for each group velocity.")
    return group_velocity

def cli(argv=None):
    parser = argparse.ArgumentParser(prog="phkit")
    parser.add_argument("--verbose", "-v", action="count", default=0, help="Increase output verbosity.")
    parser.add_argument("--exec-phonopy", default=PHONOKIT_PHONO2, help=f"Executable phonopy (default: {PHONOKIT_PHONO2})")
    parser.add_argument("--exec-vaspkit", default=PHONOKIT_VASPKIT_EXEC, help=f"Executable vaspkit (default: {PHONOKIT_VASPKIT_EXEC})")
    parser.add_argument("-a", "--abort-on-exist", action="store_true", help="Abort if output files exist")

    # subparsers = parser.add_subparsers(title="Tips", metavar="Subcommand", help="Description", dest="command")
    subparsers = parser.add_subparsers(title="subcommands", dest="command")

    parser_supercell = subparsers.add_parser("supercell", help="Create supercells with displacements.")
    parser_supercell.set_defaults(func=_options_supercell(parser_supercell))

    parser_prepare = subparsers.add_parser("prepare", help="Prepare VASP input files for IFC calculations.")
    parser_prepare.set_defaults(func=_options_prepare(parser_prepare))

    parser_collect = subparsers.add_parser("collect", help="Collect IFC results from FD or DFPT workflows.")
    parser_collect.set_defaults(func=_options_collect(parser_collect))

    parser_band = subparsers.add_parser("band", help="Calculate phonon band structure.")
    parser_band.set_defaults(func=_options_band(parser_band))

    parser_tprop = subparsers.add_parser("tprop", help="Calculate thermal properties.")
    parser_tprop.set_defaults(func=_options_tprop(parser_tprop))

    parser_vg = subparsers.add_parser("vg", help="Calculate group velocities.")
    parser_vg.set_defaults(func=_options_vg(parser_vg))

    # Enable by:
    #   $ activate-global-python-argcomplete
    #   $ eval "$(register-python-argcomplete phkit)"
    argcomplete.autocomplete(parser)
    args = parser.parse_args(argv)

    if not args.command:
        return parser.print_help()

    logging.basicConfig(level=max(logging.WARNING - args.verbose * 10, 0),
                        format="[PHONOKIT] %(message)s")

    logging.info("Starting...")
    args.func(args)
    logging.info("Done.")
